{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangChain\n",
    "\n",
    "<div class=\"subtitle\">Leverage your LangChain stack with LMQL</div>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LMQL can also be used together with the [🦜🔗 LangChain](https://python.langchain.com/en/latest/index.html#) python library. Both, using langchain libraries from LMQL code and using LMQL queries as part of chains are supported."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using LangChain from LMQL"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first consider the case, where one may want to use LangChain modules as part of an LMQL program. In this example, we want to leverage the LangChain `Chroma` retrieval model, to enable question answering about a text document (the LMQL paper in this case).\n",
    "\n",
    "First, we need to import the required libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#notebooks.js:hidden\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "nbsphinx": "hidden",
    "tags": []
   },
   "outputs": [],
   "source": [
    "#notebooks.js:hidden\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys \n",
    "sys.path.append(\"../../../../src/\")\n",
    "# load and set OPENAI_API_KEY\n",
    "import os \n",
    "from lmql.runtime.openai_secret import openai_secret\n",
    "os.environ[\"OPENAI_API_KEY\"] = openai_secret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lmql\n",
    "import asyncio\n",
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.vectorstores import Chroma"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we load and embed the text of the relevant document (`lmql.txt` in our case)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load text of LMQL paper\n",
    "with open(\"lmql.txt\") as f:\n",
    "    contents = f.read()\n",
    "texts = []\n",
    "for i in range(0, len(contents), 120):\n",
    "    texts.append(contents[i:i+120])\n",
    "\n",
    "embeddings = OpenAIEmbeddings()\n",
    "docsearch = Chroma.from_texts(texts, embeddings, \n",
    "    metadatas=[{\"text\": t} for t in texts], persist_directory=\"lmql-index\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then construct a chatbot function, using a simple LMQL query, that first prompts the user for a question via `await input(...)`, retrieves relevant text paragraphs using LangChain and then produces an answer using `openai/gpt-3.5-turbo` (ChatGPT)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#notebooks.js:show_stdout=false\n",
    "import termcolor\n",
    "\n",
    "@lmql.query(model=\"openai/gpt-3.5-turbo\")\n",
    "async def chatbot():\n",
    "    '''lmql\n",
    "    # system instruction\n",
    "    \"\"\"{:system} You are a chatbot that helps users answer questions.\n",
    "    You are first provided with the question and relevant information.\"\"\"\n",
    "    \n",
    "    # chat loop\n",
    "    while True:\n",
    "        # process user input\n",
    "        q = await input(\"\\nQuestion: \")\n",
    "        if q == \"exit\": break\n",
    "        # expose question to model\n",
    "        \"{:user} {q}\\n\"\n",
    "        \n",
    "        # insert retrieval results\n",
    "        print(termcolor.colored(\"Reading relevant pages...\", \"green\"))\n",
    "        results = set([d.page_content for d in docsearch.similarity_search(q, 4)])\n",
    "        information = \"\\n\\n\".join([\"...\" + r + \"...\" for r in list(results)])\n",
    "        \"{:system} \\nRelevant Information: {information}\\n\"\n",
    "        \n",
    "        # generate model response\n",
    "        \"{:assistant} [RESPONSE]\"\n",
    "    '''\n",
    "\n",
    "await chatbot(output_writer=lmql.stream(variable=\"RESPONSE\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```promptdown\n",
    "# Chat Log\n",
    "[bubble:system| You are a chatbot that helps users answer questions. \n",
    "You are first provided with the question and relevant information.]\n",
    "[bubble:user| What is LMQL?]\n",
    "[bubble:system| Relevant Information: (inserted by retriever)]\n",
    "[bubble:assistant| LMQL is a high-level query language for LMs that allows for great expressiveness and supports scripted prompting.]\n",
    "[bubble:user| How to write prompts?]\n",
    "[bubble:system| Relevant Information: (inserted by retriever)]\n",
    "[bubble:assistant| To write prompts, you can use a language model to expand the prompt and obtain the answer to a specific question.]\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown in the query, inline LMQL code appearing in a Python script can access the outer scope containing e.g. the `docsearch` variable, and access any relevant utility functions and object provided by LangChain.\n",
    "\n",
    "For more details on building Chat applications with LMQL, see the [Chat API documentation](../chat.md)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using LMQL from LangChain"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to using langchain utilities in LMQL query code, LMQL queries can also seamlessly be integrated as a `langchain` `Chain` component. \n",
    "\n",
    "For this consider, the sequential prompting example from the `langchain` documentation, where we first prompt the language model to propose a company name for a given product, and then ask it for a catchphrase.\n",
    "\n",
    "To get started, we first import the relevant langchain components, as well as LMQL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import LLMChain, PromptTemplate\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.prompts.chat import (ChatPromptTemplate,HumanMessagePromptTemplate)\n",
    "from langchain.llms import OpenAI\n",
    "\n",
    "import lmql"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our chain has two stages: (1) Asking the model for a company name, and (2) asking the model for a catchphrase. For the sake of this example, we will implement (1) in with a langchain prompt and (2) with an LMQL query. \n",
    "\n",
    "First, we define the langchain prompt for the company name and instantiate the resulting `LLMChain`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup the LM to be used by langchain\n",
    "llm = OpenAI(temperature=0.9)\n",
    "\n",
    "human_message_prompt = HumanMessagePromptTemplate(\n",
    "        prompt=PromptTemplate(\n",
    "            template=\"What is a good name for a company that makes {product}?\",\n",
    "            input_variables=[\"product\"],\n",
    "        )\n",
    "    )\n",
    "chat_prompt_template = ChatPromptTemplate.from_messages([human_message_prompt])\n",
    "chat = ChatOpenAI(temperature=0.9)\n",
    "chain = LLMChain(llm=chat, prompt=chat_prompt_template)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This can already be executed to produce a company name:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'VibrantSock Co.\\nColorSplash Socks\\nRainbowThreads\\nChromaSock Co.\\nKaleidosocks\\nColorPop Socks\\nPrismStep\\nSockMosaic\\nHueTrend Socks\\nSpectrumStitch\\nColorBurst Socks'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.run(\"colorful socks\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define prompt (2) in LMQL, i.e. the LMQL query generating the catchphrase:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "@lmql.query(model=\"chatgpt\")\n",
    "async def write_catch_phrase(company_name: str):\n",
    "    '''\n",
    "    \"Write a catchphrase for the following company: {company_name}. [catchphrase]\"\n",
    "    '''"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we can run this part in isolation, like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' \"Step up your style with Socks Inc. - where comfort meets fashion!\"'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(await write_catch_phrase(\"Socks Inc\")).variables[\"catchphrase\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To chain the two prompts together, we can use a `SimpleSequentialChain` from `langchain`. To make an LMQL query compatible for use with `langchain`, just call `.aschain()` on it, before passing it to the `SimpleSequentialChain` constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import SimpleSequentialChain\n",
    "overall_chain = SimpleSequentialChain(chains=[chain, write_catch_phrase.aschain()], verbose=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can run the overall chain, relying both on LMQL and langchain components:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new SimpleSequentialChain chain...\u001b[0m\n",
      "\u001b[36;1m\u001b[1;3mSockSplash\u001b[0m\n",
      "\u001b[33;1m\u001b[1;3m \"Step into freshness with SockSplash!\"\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      " \"Step into freshness with SockSplash!\"\n"
     ]
    }
   ],
   "source": [
    "#notebooks.js:show_stdout=false\n",
    "# Run the chain specifying only the input variable for the first chain.\n",
    "catchphrase = overall_chain.run(\"colorful socks\")\n",
    "print(catchphrase) "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "> Entering new SimpleSequentialChain chain...\n",
    "RainbowSocks Co.\n",
    " \"Step into a world of color with RainbowSocks Co.!\"\n",
    "\n",
    "> Finished chain.\n",
    " \"Step into a world of color with RainbowSocks Co.!\"\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Overall, we thus have a chain that combines langchain and LMQL components, and can be used as a single unit.\n",
    "\n",
    "::: info Asynchronous Use\n",
    "You may encounter problems because of the mismatch of LangChain's synchronous APIs with LMQL's `async`-first design.\n",
    "\n",
    "To avoid problems with this, you can install the [`nest_asyncio`](https://pypi.org/project/nest-asyncio/) package and call `nest_asyncio.apply()` to enable nested event loops. LMQL will then handle event loop nesting and sync-to-async conversion for you.\n",
    ":::"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
